#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :q2_v4.py
# @Time      :2024/9/7 16:20
# @Author    :YKW
# !/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :q3.py
# @Time      :2024/9/7 2:22
# @Author    :YKW
import math
import sys
import time
import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import quad
from shapely.geometry import MultiPolygon
from shapely.geometry import Polygon
import pandas as pd

p = 0.55  # 螺距，单位：米
b = p / (2 * np.pi)  # b的值
r_0_1 = 0.55 * 16  # 初始半径 r_0^(1)
theta_0_1 = 16 * 2 * np.pi  # 初始角度 theta_0^(1)
timespace = 0.1


def secant_method(f, x0, x1, now, tolerance=1e-6, max_iterations=1000):
    for i in range(max_iterations):
        fx0 = f(x0, now)
        fx1 = f(x1, now)
        if fx1 - fx0 == 0:
            return x1
        denominator = fx1 - fx0
        x2 = x1 - fx1 * ((x1 - x0) / denominator)
        if abs(x2 - x1) < tolerance:
            return x2
        x0, x1 = x1, x2
    return x2


def binary_search(f, a, b, now, tol=1e-6, max_iter=1000):
    if f(a, now) * f(b, now) >= 0:
        raise ValueError("f(a) 和 f(b) 必须异号")

    for _ in range(max_iter):
        mid = (a + b) / 2
        f_mid = f(mid, now)
        if abs(f_mid) < tol:
            return mid
        elif f(a, now) * f_mid < 0:
            b = mid
        else:
            a = mid
    print("answer not find")
    return None


def polar_distance(theta1, theta2):
    r1 = b * theta1
    r2 = b * theta2
    # 计算角度差
    angle_diff = theta2 - theta1
    # 计算距离
    distance = math.sqrt((r1 - r2 * math.cos(angle_diff)) ** 2 + (r2 * math.sin(angle_diff)) ** 2)
    return distance


def cal_body(theta, now):
    length = 1.65
    distance = polar_distance(theta, now)
    return distance - length


def cal_head(theta, now):
    length = 2.86
    distance = polar_distance(theta, now)
    return distance - length


def get_theta_position(now):  # 给龙头，给全部
    result = [now]
    num = 223 - 1 - 1
    # body = binary_search(cal_head, now + 0.001, now + 1, now)
    body = secant_method(cal_head, now + 0.001, now + 0.3, now)
    result.append(body)
    for i in range(num):
        # next = binary_search(cal_body, body + 0.001, body + 1, body)
        next = secant_method(cal_body, body + 0.001, body + 0.3, body)
        result.append(next)
        body = next
    return result


def integrand(theta):
    return np.sqrt((b * theta) ** 2 + b ** 2)


def find_theta_1(theta_1, theta_0, length=0.01):  # length=v_head*timespace
    length = timespace
    return quad(integrand, theta_1, theta_0)[0] - length


def get_theta_time(now, time_now, time_next, timespace):  # 给一点，算时序
    result = [now]
    t = 0
    while abs(t - (time_next - time_now)) > 0.01 * timespace:
        t += timespace
        # theta_prev = result[-1]
        # theta_new = fsolve(find_theta_1, theta_prev - 1, args=(theta_prev))[0]
        theta_new = secant_method(find_theta_1, now + 0.001, now + 0.3, now)
        result.append(theta_new)
        now = theta_new
    return result


def get_init_position(now, time, timespace):
    result = [now]
    t = 0
    mul = int(1 / timespace)
    idx = [i * mul for i in range(int(time))]
    # idx=[0]
    # count=0
    while abs(t - time) > 0.01 * timespace:
        t += timespace
        # count+=1
        # if abs(t-int(t))<0.01:
        #     idx.append(count)
        # theta_new = fsolve(find_theta_1, theta_prev - 1, args=(theta_prev))[0]
        theta_new = secant_method(find_theta_1, now + 0.001, now + 0.3, now)
        result.append(theta_new)
        now = theta_new
    # if len(idx)>301:
    #     idx=idx[:301]
    # result = [result[i] for i in idx]
    # print(len(result))

    return result


def cal_total():
    head = get_theta_time(theta_0_1, 0, 300, 0.01)
    ans = []
    for i in head:
        tem = get_theta_position(i)
        ans.append(tem)
    return ans


def cal_velocity_components(theta_a, theta_b, v_a):
    """计算速度分解到ab线段方向的分量"""
    # 计算螺线在a点和b点的切线角度
    r1 = b * theta_a
    r2 = b * theta_b
    x_0 = r1 * np.cos(theta_a)
    y_0 = r1 * np.sin(theta_a)
    x_b = r2 * np.cos(theta_b)
    y_b = r2 * np.sin(theta_b)
    tangent_angle_a = theta_a + np.pi / 2
    tangent_angle_b = theta_b + np.pi / 2

    # 计算ab线段的角度
    ab_angle = np.arctan2(y_b - y_0, x_b - x_0)

    # 计算速度分解到ab线段方向的分量
    v_a_line = v_a * np.cos(tangent_angle_a - ab_angle)
    v_b_line = v_a_line
    v_b = v_b_line / np.cos(tangent_angle_b - ab_angle)
    return v_b


def cal_line(point1, point2):
    x1, y1 = point1
    x2, y2 = point2

    A = y2 - y1
    B = x1 - x2

    # 计算 C
    C = - (A * x1 + B * y1)

    return (A, B, C)


def cal_points_on_line(point1, point2, length):
    l = length
    x1, y1 = point1
    x2, y2 = point2
    if x1 < x2:
        x1, x2 = x2, x1
        y1, y2 = y2, y1
    elif x1 == x2 and y1 < y2:
        x1, x2 = x2, x1
        y1, y2 = y2, y1
    # 计算直线的方向向量
    direction = [x2 - x1, y2 - y1]

    # 计算方向向量的模长
    direction_length = np.sqrt(direction[0] ** 2 + direction[1] ** 2)

    # 标准化方向向量
    direction_unit = [direction[0] / direction_length, direction[1] / direction_length]

    # 计算两个方向上的点
    if direction_unit[0] < 0 or (direction_unit[0] == 0 and direction_unit[1] < 0):
        direction_unit[0], direction_unit[1] = -direction_unit[0], -direction_unit[1]
    point1_plus = [x1 + direction_unit[0] * l, y1 + direction_unit[1] * l]
    point1_minus = [x2 - direction_unit[0] * l, y2 - direction_unit[1] * l]

    return point1_plus, point1_minus


def cal_vertical_dot(point1, point2, length):
    l = length
    x1, y1 = point1
    x2, y2 = point2
    direction = [x2 - x1, y2 - y1]
    direction_length = np.sqrt(direction[0] ** 2 + direction[1] ** 2)
    direction_unit = [direction[0] / direction_length, direction[1] / direction_length]
    direction_unit = [-direction_unit[1], direction_unit[0]]
    if direction_unit[0] < 0 or (direction_unit[0] == 0 and direction_unit[1] < 0):
        direction_unit[0], direction_unit[1] = -direction_unit[0], -direction_unit[1]
    point1_plus = [x1 + direction_unit[0] * l, y1 + direction_unit[1] * l]
    point1_minus = [x2 - direction_unit[0] * l, y2 - direction_unit[1] * l]
    return point1_plus, point1_minus


class node:
    p = 0.55  # 螺距，单位：米
    b = p / (2 * np.pi)  # b的值

    def __init__(self, time, position, speed=None, theta=None):
        self.length = 2.86 if position == 0 else 1.65
        self.type = 0 if position == 0 else 1
        self.theta = theta
        self.r = None
        self.x = None
        self.y = None
        self.speed = speed
        self.speed_x = None
        self.speed_y = None
        self.time = time
        self.before = None
        self.after = None
        self.box = None

    def refresh(self):
        if self.theta is None:
            print("ERROR NO Theta")
            return 0
        self.r = b * self.theta
        self.x = self.r * np.cos(self.theta)
        self.y = self.r * np.sin(self.theta)
        if self.before is not None:
            self.speed = cal_velocity_components(self.before.theta, self.theta, self.before.speed)


class bench:
    def __init__(self, node1, node2):
        p1 = (node1.x, node1.y)
        p2 = (node2.x, node2.y)
        mid1, mid2 = cal_points_on_line(p1, p2, length=0.275)
        self.dotA, self.dotB = cal_vertical_dot(mid1, mid2, length=0.15)
        self.dotC, self.dotD = cal_vertical_dot(mid2, mid1, length=0.15)
        shell = [self.dotA, self.dotD, self.dotB, self.dotC]
        self.box = Polygon(shell)


class Dragon:
    start_theat = 16 * 2 * np.pi

    def __init__(self, time=0, timespace=1, body_num=222, head_speed=1):
        self.time = time
        self.timespace = timespace
        time_list = get_init_position(self.start_theat, self.time, self.timespace)
        self.head = node(time, 0, speed=head_speed, theta=time_list[-1] if time else self.start_theat)
        pos_list = get_theta_position(time_list[-1])
        self.length = body_num + 1
        self.body = [node(time, position, theta=pos_list[position])
                     for position in range(1, body_num + 1)]

        self.head.after = self.body[0]
        self.body[0].after = self.body[1]
        for i in range(1, len(self.body) - 1):
            self.body[i].before = self.body[i - 1]
            self.body[i].after = self.body[i + 1]
        self.body[0].before = self.head
        self.body[-1].before = self.body[-2]
        self.state = 1

    def refresh_total(self):
        self.head.refresh()
        for i in self.body:
            i.refresh()

    def check_state(self):  # 1 if 未翻转 else 0
        self.state = 0

    def __move_body__(self):
        pos_list = get_theta_position(self.head.theta)
        for i in range(len(self.body)):
            self.body[i].theta = pos_list[i + 1]
            self.body[i].refresh()

    def __move_head__(self, time):
        time_list = get_theta_time(self.head.theta, self.time, time, self.timespace)
        self.time = time
        self.head.theta = time_list[-1]
        self.head.refresh()

    def time_check(self, time):
        if self.state:
            self.__move_head__(time)
            self.__move_body__()

    def move_forward(self):
        if self.state:
            pass

    def __box_refresh__(self):
        self.head.box = bench(self.head, self.head.after).box
        for i in self.body:
            if i.after is not None:
                i.box = bench(i, i.after).box

    def check_crash(self):
        self.__box_refresh__()
        # print(self.time)
        boxes = [self.head.box]
        # boxes=boxes[:220]
        for i in self.body:
            if i.box is not None:
                boxes.append(i.box)

        for i in range(2, len(boxes)):
            flag = boxes[0].intersects(boxes[i])
            if flag:
                return (0, i)
        return -1, -1

        # for i in range(len(boxes)):
        #     for j in range(len(boxes)):
        #         if abs(i - j) <= 1:
        #             continue
        #         flag = boxes[i].intersects(boxes[j])
        #         if flag:
        #             return (i, j)
        # return -1, -1

    def draw(self):
        self.__box_refresh__()
        boxes = [self.head.box]
        # boxes=boxes[:220]
        for i in self.body:
            if i.box is not None:
                boxes.append(i.box)
        multipolygon = MultiPolygon(boxes)
        # 创建一个图形和轴
        fig, ax = plt.subplots()

        # 遍历 MultiPolygon 中的每个 Polygon 并绘制它们
        for polygon in multipolygon.geoms:
            x, y = polygon.exterior.xy
            ax.plot(x, y)

        # 设置坐标轴的比例相同
        ax.set_aspect('equal')

        # 显示图形
        plt.show()


def transpose_2d(data):
    transposed = list(map(list, zip(*data)))
    return transposed


def question1():
    with pd.ExcelWriter('result1.xlsx') as writer:
        # 算极角
        # head = get_theta_time(theta_0_1, 0, 300, 0.01)  # 300s的龙头极角变化
        theta_all = cal_total()  # 每时刻，每个点的极角
        print("完成极角计算")

        # 算坐标
        pos = []
        for i in theta_all:
            tem = []
            for j in i:
                r = b * j
                x = r * np.cos(j)
                y = r * np.sin(j)
                tem.extend([x, y])
            pos.append(tem)
        df = pd.DataFrame(pos).transpose()
        index = ["龙头x (m)", "龙头y (m)"]
        columns = [f"{i} s" for i in range(301)]
        for i in range(1, 222):
            index.extend([f"第{i}节龙身x (m)", f"第{i}节龙身y (m)"])
        index.extend(["龙尾x (m)", "龙尾y (m)", "龙尾（后）y (m)", "龙尾（后）y (m)"])
        df.columns = columns
        df.index = index
        df.to_excel(writer, sheet_name="坐标", engine="openpyxl")
        print("完成坐标计算")
        # 算速度
        vel = [[1] for i in range(0, 301)]
        for time in range(0, 301):  # [0,300]时间
            for point in range(1, 224):  # [0,223]个点
                theta_a = theta_all[time][point - 1]
                theta_b = theta_all[time][point]
                v_a = vel[time][point - 1]
                if theta_b == 0:
                    vel[time].append(0)
                else:
                    vel[time].append(cal_velocity_components(theta_a, theta_b, v_a))
        vel = transpose_2d(vel)
        index = ["龙头 (m/s)"]
        for i in range(1, 222):
            index.append(f"第{i}节龙身 (m/s)")
        index.extend(["龙尾 (m/s)", "龙尾（后） (m/s)"])
        df = pd.DataFrame(vel, index=index, columns=columns)
        df.to_excel(writer, sheet_name="速度", engine="openpyxl")
        print("完成速度计算")


def question2():
    timespace = 0.1
    times = 300
    dragon = Dragon(time=times, timespace=timespace)
    tem = []
    teem = []
    count = 0
    eval_fre = 10
    # for times in range(290, 350):
    #     dragon.time_check(times)
    #     dragon.draw()
    #     print(dragon.head.theta)

    while times < 450:
        times += timespace
        count += 1
        # if count%eval_fre==0:
        #     dragon.draw()
        dragon.time_check(times)
        i, j = dragon.check_crash()
        # dragon.draw()
        if i != -1:
            # dragon.draw()
            print("Crash!!!")
            print(i, j)
            tem.append((i, j))
            teem.append(times)
            break
        else:
            print(dragon.head.theta)


if __name__ == "__main__":
    timespace = 0.01
    times = 410
    dragon = Dragon(time=times, timespace=timespace)
    tem = []
    teem = []
    count = 0
    eval_fre = 10
    # for times in range(290, 350):
    #     dragon.time_check(times)
    #     dragon.draw()
    #     print(dragon.head.theta)

    while times < 450:
        times += timespace
        count += 1
        # if times> 412.:
        #     dragon.draw()
        dragon.time_check(times)
        i, j = dragon.check_crash()
        # dragon.draw()
        if i != -1:
            dragon.draw()
            # plt.savefig("C:\\Users\\Administrator\\Desktop\\",fname='scatter.png',dpi=600)
            print("Crash!!!")
            print(i, j)
            tem.append((i, j))
            teem.append(times)
            break
        else:
            print(times)
            print(dragon.head.theta)
